"""
Helper module for traversing and querying the ARB and NCBI taxonomy databases.
"""

from ssummo.traverse import dict_walk
from ssummo import CONFIG

import os
import re
import sys
from threading import Thread

class TaxDB:
    def __init__(self):
        import MySQLdb
        self.user = CONFIG.user
        self.host = CONFIG.host
        self.passwd = CONFIG.passwd
        self.db = CONFIG.db
        self.cnx = MySQLdb.connect(user=self.user, host=self.host,
                                   passwd=self.passwd, db=self.db)
        self.cur = self.cnx.cursor()

    def fetch_tax_id(self, table, OTUName, parentName, Queue=None):
        self.cur.execute('SELECT tax_ID,rank FROM {0} WHERE Name="{1}" AND '
                         'ParentName="{2}";'.format(table, OTUName, parentName))
        IDs = set( (str(r[0]), r[1] ) for r in self.cur.fetchall() )
        if len(IDs) > 1:
            print('Multiple IDs!!')
            print('table:\t{0}\nOTU\t{1}\nparent:\t{2}\nID:\t{3}'
                  .format(table, OTUName, parentName, repr(IDs)))
            retVal = (0, 'unknown')
        elif len(IDs) == 0:
            self.cur.execute('SELECT tax_ID,rank FROM {0} WHERE Name="{1}";'
                             .format(table, OTUName))
            results = set((str(tax_ID[0]), tax_ID[1])
                          for tax_ID in self.cur.fetchall() )
            if len(results) == 1:
                retVal = results.pop()
            elif len(results) == 0:
                if ' ' in OTUName:
                    return self.fetch_tax_id(table,
                                             OTUName.rsplit(' ', 1)[0].strip(),
                                             parentName, Queue=Queue)
                retVal = (0, 'unknown')
                pass
            else:
                print("Selecting Name '{0}' from '{1}' gives these results:-"
                      .format(OTUName, table))
                print("{0}".format(results))
                retVal = (0, 'unknown')
        else:
            retVal = IDs.pop()
        sys.stdout.flush()
        if Queue is not None:
            Queue.put(retVal)
        else:
            return retVal

    #def TaxIDThread(self, inQueue, outPipe):
    def _tax_ID_thread(self, in_queue, out_pipe):
        """Fetch taxonomic IDs from the MySQL database asynchronously. Designed
        to be run in a separate thread. (:sa:`.self.create_thread`)

        :param in_queue: :class:`Queue.Queue` instance, from which to get tuples
                         of the form (table, OTU, parent)
        :param out_pipe: :class:`multiprocessing.Pipe` instance, to which
                         results of the form ``{ OTU : ('tax_ID', 'rank')}`` are
                         sent.
        """
        x = in_queue.get()
        while x != 'STOP':
            if len(x) == 3:
                table, OTU, parent = x
            else:
                print(x)
                err = "Items put into this queue should be of the form: "
                err += "(table, OTUname, parentName).\n"
                err += "Not {0}".format(x)
                raise ValueError(err)
            out_pipe.send({ OTU : self.fetch_tax_id(table, OTU, parent)})
            x = in_queue.get()

    def create_thread(self, in_queue, out_pipe):
        """Create and return a :class:`threading.Thread` instance, meant for
        running :func:`.self._tax_ID_thread` in a separate thread.

        :sa: :func:`TaxDB._tax_ID_thread`
        """
        return Thread(target=self._tax_ID_thread, group=None,
                      name="ssummo.taxonomy.TaxDB",
                      args=(in_queue, out_pipe), kwargs=None)

    def get_ranks(self , table='Prokaryotes'):
        self.cur.execute( 'SELECT DISTINCT(Rank) FROM {0};'.format( table ) )
        ranks = [ rank[0] for rank in self.cur.fetchall() ]
        return ranks

    def __del__(self):
            self.cnx.close()

def fetch_rank(TaxDBObj, table, OTUName, parent_name=None):
    """
    Checks MySQL database for taxon with name OTUName, and
    parent parent_name.
    If a unique rank is found, return that, otherwise return
    False.
    """
    TaxDBObj.cur.execute('SELECT Rank,Name from {0} where Name="{1}";'
                         .format(table, OTUName) )
    rank = set( r for r in TaxDBObj.cur.fetchall()  )
    if len(rank) == 1:
        return rank.pop()
    elif len(rank) == 0:
        TaxDBObj.cur.execute('SELECT Rank,Name from {0} where Name LIKE "{1}%";'
                             .format(table, OTUName) )
        rank = set( r for r in TaxDBObj.cur.fetchall()  )
        if len(rank) == 1:
            return rank.pop()
        elif len(rank) == 0:
            next_name = OTUName.rsplit(' ', 1)
            if len( next_name ) == 2:
                return fetch_rank(TaxDBObj, table, next_name[0],
                                  parent_name=parent_name )
        else:
            if parent_name is not None:
                TaxDBObj.cur.execute('SELECT Rank, Name from {0} where '
                                     'Name="{1}" and ParentName="{2}";'
                                     .format(table, OTUName, parent_name))
                rank = set( r for r in TaxDBObj.cur.fetchall() )
                if len( rank ) == 1:
                    return rank.pop()
    else:  # more than one rank.
        if parent_name is not None:
            TaxDBObj.cur.execute('SELECT Rank, Name from {0} where Name="{1}" '
                                 'and ParentName="{2}";'
                                 .format(table, OTUName, parent_name ) )
            rank = set( r for r in TaxDBObj.cur.fetchall() )
            if len( rank ) == 1:
                return rank.pop()
    return (False, False)

def find_taxa(tdict, options):
    found = []
    interesting = re.compile(r'|'.join( options['-in'] ), re.I)
    sep = os.path.sep
    for path, node in dict_walk('', tdict):
        if interesting.search( path.rsplit(sep, 1)[-1] ):
            found.append(path)
    return found

def get_ranks(taxPathList, TaxDBObj):
    """Give a list of all taxonomies you've got, and a TaxDB object.
    This will return a list of all the ranks, mirroring the taxPathList
    given, and also the table name.
    """
    ranks = ['root']
    taxPath = ','.join(taxPathList)
    domMatch = re.search(r'(Bacteria)|(Eukaryota)|(Archaea)', taxPath, re.I)
    if not domMatch:
        ProRanks = []
        for i in xrange(1, len( taxPathList ) ):
            # Appending ranks.
            ProRanks.append(str( TaxDBObj.fetch_tax_id(
                            'Prokaryotes', taxPathList[i],
                            taxPathList[i-1] )[1] ).lower() )
        EukRanks = []
        for i in xrange(1, len( taxPathList ) ):
            EukRanks.append(str(TaxDBObj.fetch_tax_id(
                            'Eukaryotes', taxPathList[i],
                            taxPathList[i-1] )[1] ).lower())
        if EukRanks.count( 'unknown') > ProRanks.count( 'unknown' ):
            table = 'Prokaryotes'
            ranks = ProRanks
        else:
            table = 'Eukaryotes'
            ranks = EukRanks
    else:
        if domMatch.groups()[1] is not None:
            table = 'Eukaryotes'
        else:
            table = 'Prokaryotes'
        for i in xrange(1, len( taxPathList ) ):
            ranks.append(str(TaxDBObj.fetch_tax_id(table, taxPathList[i],
                         taxPathList[i-1] )[1]).lower())
    return ranks, table

def match_ranks(taxa):
    TaxDBObj = TaxDB()
    matchedRanks = {}
    sep = os.path.sep
    for taxon in taxa:
        ranks, table = get_ranks( taxon.split( sep ), TaxDBObj )
        matchedRanks.update( { taxon : ranks } )
    return matchedRanks

def walk_ranks(tdict, path=''):
    names = tdict.keys()
    del( names[names.index('rank')] )
    yield (names, tdict, path, )
    join = os.path.join
    for name in names:
        for x in walk_ranks(tdict[name], path=join(path, name) ):
            yield x
